{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "cfb64210-9c6b-47d7-81f4-67dbdab68e4c",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Composable Graph with Weaviate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa0e62b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "import weaviate\n",
    "from pprint import pprint\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
    "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))\n",
    "\n",
    "from llama_index import (\n",
    "    VectorStoreIndex,\n",
    "    SimpleKeywordTableIndex,\n",
    "    ListIndex,\n",
    "    VectorStoreIndex,\n",
    "    SimpleDirectoryReader,\n",
    ")\n",
    "from llama_index.vector_stores import WeaviateVectorStore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b594b69-5814-4ff1-abc0-765b724f6339",
   "metadata": {},
   "outputs": [],
   "source": [
    "resource_owner_config = weaviate.AuthClientPassword(\n",
    "    username=\"<username>\",\n",
    "    password=\"<password>\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d6297f6-1a78-4dc5-9d48-f3968729e273",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = weaviate.Client(\n",
    "    \"https://test-weaviate-cluster.semi.network/\",\n",
    "    auth_client_secret=resource_owner_config,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5583b867-ab33-4e0e-8a38-9995615faa84",
   "metadata": {},
   "outputs": [],
   "source": [
    "# [optional] set batch\n",
    "client.batch.configure(batch_size=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49e0d841-680f-4a0c-b455-788b54978ebf",
   "metadata": {},
   "source": [
    "#### Load Datasets\n",
    "\n",
    "Load both the NYC Wikipedia page as well as Paul Graham's \"What I Worked On\" essay"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ec16a8b-6aae-4bf7-9b83-b82087b4ea52",
   "metadata": {},
   "outputs": [],
   "source": [
    "# fetch \"New York City\" page from Wikipedia\n",
    "from pathlib import Path\n",
    "\n",
    "import requests\n",
    "\n",
    "response = requests.get(\n",
    "    \"https://en.wikipedia.org/w/api.php\",\n",
    "    params={\n",
    "        \"action\": \"query\",\n",
    "        \"format\": \"json\",\n",
    "        \"titles\": \"New York City\",\n",
    "        \"prop\": \"extracts\",\n",
    "        # 'exintro': True,\n",
    "        \"explaintext\": True,\n",
    "    },\n",
    ").json()\n",
    "page = next(iter(response[\"query\"][\"pages\"].values()))\n",
    "nyc_text = page[\"extract\"]\n",
    "\n",
    "data_path = Path(\"data\")\n",
    "if not data_path.exists():\n",
    "    Path.mkdir(data_path)\n",
    "\n",
    "with open(\"../test_wiki/data/nyc_text.txt\", \"w\") as fp:\n",
    "    fp.write(nyc_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39c00aeb-adef-4ce3-8134-031de18e64ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load NYC dataset\n",
    "nyc_documents = SimpleDirectoryReader(\"../test_wiki/data/\").load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddff8f98-e002-40c5-93ac-93aa40dca5ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load PG's essay\n",
    "essay_documents = SimpleDirectoryReader(\"../paul_graham_essay/data/\").load_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1782198-c0de-4679-8951-1297c21b8639",
   "metadata": {},
   "source": [
    "### Building the document indices\n",
    "Build a tree index for the NYC wiki page and PG essay"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5431e83e-428b-4473-bad1-24b7a6c4db38",
   "metadata": {},
   "outputs": [],
   "source": [
    "# build NYC index\n",
    "from llama_index.storage.storage_context import StorageContext\n",
    "\n",
    "\n",
    "vector_store = WeaviateVectorStore(weaviate_client=client, class_prefix=\"Nyc_docs\")\n",
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "nyc_index = VectorStoreIndex.from_documents(\n",
    "    nyc_documents, storage_context=storage_context\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b5aad4a-49ef-4b24-962a-0793f4f09316",
   "metadata": {},
   "outputs": [],
   "source": [
    "# build essay index\n",
    "vector_store = WeaviateVectorStore(weaviate_client=client, class_prefix=\"Essay_docs\")\n",
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "essay_index = VectorStoreIndex.from_documents(\n",
    "    essay_documents, storage_context=storage_context\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdcb22d5-4df8-4d65-aa29-6493fc027fe2",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Set summaries for the indices\n",
    "\n",
    "Add text summaries to indices, so we can compose other indices on top of it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4149cbbd-7d0b-48c4-8c47-7d67ae0c55f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "nyc_index_summary = \"\"\"\n",
    "    New York, often called New York City or NYC, \n",
    "    is the most populous city in the United States. \n",
    "    With a 2020 population of 8,804,190 distributed over 300.46 square miles (778.2 km2), \n",
    "    New York City is also the most densely populated major city in the United States, \n",
    "    and is more than twice as populous as second-place Los Angeles. \n",
    "    New York City lies at the southern tip of New York State, and \n",
    "    constitutes the geographical and demographic center of both the \n",
    "    Northeast megalopolis and the New York metropolitan area, the \n",
    "    largest metropolitan area in the world by urban landmass.[8] With over \n",
    "    20.1 million people in its metropolitan statistical area and 23.5 million \n",
    "    in its combined statistical area as of 2020, New York is one of the world's \n",
    "    most populous megacities, and over 58 million people live within 250 mi (400 km) of \n",
    "    the city. New York City is a global cultural, financial, and media center with \n",
    "    a significant influence on commerce, health care and life sciences, entertainment, \n",
    "    research, technology, education, politics, tourism, dining, art, fashion, and sports. \n",
    "    Home to the headquarters of the United Nations, \n",
    "    New York is an important center for international diplomacy,\n",
    "    an established safe haven for global investors, and is sometimes described as the capital of the world.\n",
    "\"\"\"\n",
    "essay_index_summary = \"\"\"\n",
    "    Author: Paul Graham. \n",
    "    The author grew up painting and writing essays. \n",
    "    He wrote a book on Lisp and did freelance Lisp hacking work to support himself. \n",
    "    He also became the de facto studio assistant for Idelle Weber, an early photorealist painter. \n",
    "    He eventually had the idea to start a company to put art galleries online, but the idea was unsuccessful. \n",
    "    He then had the idea to write software to build online stores, which became the basis for his successful company, Viaweb. \n",
    "    After Viaweb was acquired by Yahoo!, the author returned to painting and started writing essays online. \n",
    "    He wrote a book of essays, Hackers & Painters, and worked on spam filters. \n",
    "    He also bought a building in Cambridge to use as an office. \n",
    "    He then had the idea to start Y Combinator, an investment firm that would \n",
    "    make a larger number of smaller investments and help founders remain as CEO. \n",
    "    He and his partner Jessica Livingston ran Y Combinator and funded a batch of startups twice a year. \n",
    "    He also continued to write essays, cook for groups of friends, and explore the concept of invented vs discovered in software. \n",
    "\n",
    "\"\"\"\n",
    "index_summaries = [nyc_index_summary, essay_index_summary]\n",
    "nyc_index.set_index_id(\"nyc_index\")\n",
    "essay_index.set_index_id(\"essay_index\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4d3cd8b-4134-4cfa-8002-e0a34694d2e1",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Build Keyword Table Index on top of vector indices! \n",
    "\n",
    "We set summaries for each of the NYC and essay indices, and then compose a keyword index on top of it."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eebbc448-1e0b-402c-b37e-f93bfcc0bf4f",
   "metadata": {},
   "source": [
    "### Define Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d68750c-e5ae-481a-8b03-6173020c9bf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.indices.composability import ComposableGraph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f975514f-fddd-4737-91de-97bc61394ea9",
   "metadata": {},
   "outputs": [],
   "source": [
    "graph = ComposableGraph.from_indices(\n",
    "    SimpleKeywordTableIndex,\n",
    "    [nyc_index, essay_index],\n",
    "    index_summaries=index_summaries,\n",
    "    max_keywords_per_chunk=50,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56092d98",
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_query_engines = {\n",
    "    graph.root_id: graph.root_index.as_query_engine(retriever_mode=\"simple\")\n",
    "}\n",
    "\n",
    "query_engine = graph.as_query_engine(\n",
    "    custom_query_engines=custom_query_engines,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3c4e58b-b153-4e43-bc02-274a85babbe8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set Logging to DEBUG for more detailed outputs\n",
    "# ask it a question about NYC\n",
    "response = query_engine.query(\n",
    "    \"What is the weather of New York City like? How cold is it during the winter?\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0a43443-3e00-4e48-b3ab-f6369191d53a",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c78bc3da-6bad-4998-9a81-90a3fa9200a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get source of response\n",
    "print(response.get_formatted_sources())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b53e45e-93aa-4b49-a497-ab403f6254f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ask it a question about PG's essay\n",
    "response = query_engine.query(\n",
    "    \"What did the author do growing up, before his time at Y Combinator?\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06dc71bb-882d-49f5-8566-69b0ea5019dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0894565-2b2c-4987-a891-17ba44d775b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get source of response\n",
    "print(response.get_formatted_sources())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c103b6d-0946-48ba-a875-476c706f8560",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.11.0 ('llama')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  },
  "vscode": {
   "interpreter": {
    "hash": "775fd5332502f2902173832d699e1edc37222ebadd0e97b5c8a1a7431bebae89"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
